import os

import numpy as np
import cv2
import torch
import torchvision

class Crack(torch.utils.data.Dataset):
    def __init__(self, img_path, label_path):
        self.img_path = img_path
        self.label_path = label_path
        self.img_names = os.listdir(self.img_path)
        self.label_names = os.listdir(self.label_path)
        self.transformer = torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Resize((1250, 500)),  # 注意与cv2.resize不同，torch使用(H, W)形式，
            torchvision.transforms.Normalize(mean=[0.734, 0.734, 0.7340], std=[0.0917, 0.0917, 0.0917])
        ])

    def __getitem__(self, item):
        img = cv2.imread(os.path.join(self.img_path, self.img_names[item]))
        img = self.transformer(img)
        label = cv2.imread(os.path.join(self.label_path, self.label_names[item]), 0)
        label = cv2.resize(label, (500, 1250))
        label[label<=128] = 0
        label[label>128] = 1
        return (img, label.astype(np.int64))

    def __len__(self):
        return len(self.img_names)

if __name__ == '__main__':
    dataset = Crack(r'./data/train/imgs', r'./data/train/masks')
    dataset_loader = torch.utils.data.DataLoader(dataset = dataset, batch_size=4, shuffle=True)

    for imgs, labels in dataset_loader:
        print(imgs.shape)
        print(labels.shape)
        break